from __future__ import print_function

import glob
import random
import logging as log
import os
import pathlib
import json
import cv2
import numpy as np

import torch
import sys
import time
from tqdm import tqdm

sys.path.append("/project/train/src_repo/yolov5_js")
from utils.torch_utils import select_device
# from utils.torch_utils import select_device, load_classifier
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.general import apply_classifier
from utils.augmentations import letterbox
from models.experimental import attempt_load

if __name__ == "__main__":
    class time_count:
        def __init__(self):
            self.count = 0
            self.count1_call_times = 0
            self.names = []

        def init(self):
            self.count += 1
            self.count_call_times = 0

            if self.count == 1:
                pass
            elif self.count == 2:
                self.list = [[] for _ in range(self.count1_call_times)]
            elif self.count > 2:
                pass
            else:
                raise Exception("time_count")
            self.start = time.time()

        def __call__(self, name):
            self.end = time.time()
            if self.count == 1:
                self.count1_call_times += 1
                self.names.append(name)
            elif self.count > 1:
                self.list[self.count_call_times].append(self.end - self.start)
                self.count_call_times += 1
            else:
                raise Exception("time_count")
            self.start = time.time()

        def summury(self):
            print("==" * 20)
            spend_time_mean_overall = 0
            for idx in range(len(self.names)):
                name = self.names[idx]
                spend_time_list = self.list[idx]
                spend_time_mean = np.mean(spend_time_list) * 1000
                spend_time_std = np.std(spend_time_list) * 1000
                if spend_time_mean > 0.01:
                    print("{: <80}{: <20}{: <20}".format(name, round(spend_time_mean, 2), round(spend_time_std, 10)))
                spend_time_mean_overall += spend_time_mean
            print("{: <80}{: <20}".format("overall_time", round(spend_time_mean_overall, 1)))
            print("{: <80}{: <20}".format("overall_fps", round(1000 / spend_time_mean_overall, 1)))
            print("==" * 20)
else:
    class time_count:
        def __init__(self):
            pass

        def init(self):
            pass

        def __call__(self, name):
            pass

        def summury(self):
            pass

global tc
device = 'cuda:0'
device = select_device(device)
tc = time_count()

# change_here
# ============================================================================
label_id_map = {
    0: "slagcar",
    1: "car",
    2: "tricar",
    3: "motorbike",
    4: "bicycle",
    5: "bus",
    6: "truck",
    7: "tractor"
}
imgsz = 768
model_name = "yolov5n_s768_e60_b80"

classify = False
modelc = None
half = False
stride = None
# ============================================================================

# log.basicConfig(level=log.DEBUG)
log.basicConfig(level=log.INFO)
model_weight = "/project/train/models/yolov5/{}/weights/last.pt".format(model_name)

if __name__ == "__main__":

    model_name = "local_yolov5s_s704_e100_b64"
    show_dir = "/project/train/show_ji"
    model_weight = "/project/train/models/yolov5/{}/weights/last.pt".format(model_name)
    imgpaths = glob.glob("/home/data/newdata/*.jpg")
    # 随机选取一张
    # imgpath = imgpaths[int((len(imgpaths) - 1) * random.random())]

    if not os.path.exists(show_dir):
        os.mkdir(show_dir)
    # show_path = show_dir + "/" + os.path.basename(imgpath)
    # log.info("imgpath: " + imgpath)
    # log.info("show_path: " + show_path)


# ============================================================================

@torch.no_grad()
def init():
    # log.debug("init : start")
    # log.info("model_weight: " + model_weight)

    # if not os.path.isfile(model_weight):
    #     log.error(f'{model_weight} does not exist')
    #     return None

    # Load model
    model = attempt_load(model_weight, map_location=device)
    if half:
        model.half()  # to FP16
    # =============== second-stage classifier ================
    # if classify:
    #     global modelc
    #     modelc = load_classifier(name='resnet50', n=2)  # initialize
    #     modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval()

    # global imgsz, stride
    # stride = int(exec_net.stride.max())  # model stride
    # imgsz = check_img_size(imgsz, s=stride)
    model.eval()

    return model


def post_process(objects):
    slagcar = []
    for o in objects:
        if o["name"] == "slagcar":
            slagcar.append(o)
    res = {
        "algorithm_data": {
            "is_alert": len(slagcar) > 0,
            "target_count": len(slagcar),
            "target_info": slagcar
        },
        "model_data": {
            "objects": objects
        }
    }
    # res = {
    #     "algorithm_data": {
    #         "is_alert": len(objects) > 0,
    #         "target_count": len(objects),
    #         "target_info": objects
    #     },
    #     "model_data": {
    #         "objects": objects
    #     }
    # }

    return json.dumps(res)


def process_image(handle=None, input_image=None, args=None, **kwargs):
    if input_image is None:
        log.info("None : return a empty result")
        return post_process([])
    tc.init()
    # stride = int(exec_net.stride.max())  # model stride
    stride = 32
    # imgsz = check_img_size(imgsz, s=stride)
    img = letterbox(input_image, imgsz, stride=stride, auto=True)[0]  # BGR
    tc('letterbox')

    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    img = np.expand_dims(img, axis=0)
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img.copy()).to(device)
    # tc('ascontiguousarray')
    # img = img.half() if half else img.float()
    img = img / 255.0
    tc('255')

    # pred
    pred = handle(img, augment=False, visualize=False)[0]
    tc('exec_net')

    # NMS
    log.debug("process_image : NMS")
    conf_thres = 0.25  # confidence threshold
    iou_thres = 0.45  # NMS IOU threshold
    classes = None  # filter by class: --class 0, or --class 0 2 3
    agnostic_nms = False  # class-agnostic NMS
    max_det = 1000  # maximum detections per image
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
    tc('nms')
    # ============== Second-stage classifier (optional) ==================
    # log.debug("process_image : classify")
    # if classify:
    #     pred = apply_classifier(pred, modelc, img, input_image)

    det = pred[0]
    detect_objs = []
    if len(det):
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], input_image.shape).round()
        for proposal in det:
            x0, y0, x1, y1, conf, label = proposal
            # detect_objs.append({
            #     'xmin': int(x0),
            #     'ymin': int(y0),
            #     'xmax': int(x1),
            #     'ymax': int(y1),
            #     'confidence': float(conf),
            #     # 'name': label_id_map[int(label)]
            #     'name': "mound"
            # })
            detect_objs.append({
                'x': int(x0),
                'y': int(y0),
                'width': int(x1 - x0),
                'height': int(y1 - y0),
                'confidence': float(conf),
                'name': label_id_map[int(label)]
                # 'name': "mound"
            })
    res = post_process(detect_objs)

    return res


def show(img, result, show_path):
    result = json.loads(result)
    detect_objs = result["objects"]
    imgshow = img.copy()
    if len(detect_objs) > 0:
        for ann in detect_objs:
            xmin = int(ann["xmin"])
            ymin = int(ann["ymin"])
            xmax = int(ann["xmax"])
            ymax = int(ann["ymax"])
            confidence = float(ann["confidence"])
            name = ann["name"]
            imgshow = cv2.rectangle(imgshow, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
            imgshow = cv2.putText(imgshow, name, (int(xmin), int(ymin)), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 2)
            imgshow = cv2.putText(imgshow, str(int(confidence * 100)), (int(xmin), int(ymin)), cv2.FONT_HERSHEY_SIMPLEX,
                                  1.2, (0, 0, 255), 2)
    cv2.imwrite(show_path, imgshow)


if __name__ == '__main__':

    predictor = init()
    # img = cv2.imread(imgpath)
    for i in tqdm(imgpaths):
        img = cv2.imread(i)
        result = process_image(predictor, img)

    tc.summury()
    # show(img, result, show_path)
